import torch
import numpy as np
import uiautomator2 as u2  # 安卓自动化
from Face_recognition.retinaface import Retinaface
import cv2
from Scoring_model import Nets
from PIL import Image
import time
import torchvision.transforms as transforms


def load_model(pretrained_dict, new):
    model_dict = new.state_dict()
    # 1. filter out unnecessary keys
    pretrained_dict = {k: v for k, v in pretrained_dict['state_dict'].items() if k in model_dict}
    # 2. overwrite entries in the existing state dict
    model_dict.update(pretrained_dict)
    new.load_state_dict(model_dict)


# 保存图片
s_time = 1  # 自动化等待时间
mode = "predict"
d = u2.connect('21e582f8')  # d=device对象
for fig_num in range(1000):
    d.screenshot().save(
        r"D:\pytorch\Training_account\figure\image\fig%d.png" % fig_num)  # 将第一个封面保存，命名为1。。。。
    # 将保存的图片放入人脸识别模型中识别
    retinaface = Retinaface()
    img = './figure/image/fig' + str(fig_num) + '.png'
    image = cv2.imread(img)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    r_image, b = retinaface.detect_image(image)
    if b == 10 ^ 4:
        d.swipe_ext("up", scale=1)
        continue
    s_image = r_image[b[1]:b[3], b[0]:b[2]]  # 获取截图
    s_image = cv2.cvtColor(s_image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(r"D:\pytorch\Training_account\figure\screenshots\fig%d.png" % fig_num, s_image)  # 保存截图后的图片
    # 放入打分模型
    net = Nets.AlexNet().cuda()
    load_model(torch.load('./Scoring_model/models/alexnet.pth', encoding='ISO-8859-1'), net)
    net.eval()
    sco_img = Image.open(r"D:\pytorch\Training_account\figure\screenshots\fig%d.png" % fig_num).convert('RGB')  # sco_img为张量
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])
    sco_img = transform(sco_img)
    sco_img = sco_img.unsqueeze(0).cuda(non_blocking=True)
    output = net(sco_img).squeeze(1)
    pred = output.cpu()[0]  # tensor([],device='cuda:0')
    if pred.detach().numpy() > 3.5:  # 更改这个阈值，需要多次尝试
        # 自动化刷新、点赞、收藏
        d.click(0.935, 0.596)
        time.sleep(0.5)
        d.click(0.92, 0.752)
        time.sleep(0.5)
    d.swipe_ext("up", scale=1)
